﻿using System;
using System.Linq;
using weka.core;
using Attribute=weka.core.Attribute;
using MLSharp.Weka.Interop;

namespace MLSharp.Weka
{
	/// <summary>
	/// Provides methods for converting to/from a <see cref="IDataSet"/> to/from an 
	/// <see cref="Instances"/> collection.
	/// </summary>
	public static class DataSetConverter
	{
		/// <summary>
		/// Converts the specified data set into an equivalent 
		/// <see cref="Instances"/> object.
		/// </summary>
		/// <param name="dataSet"></param>
		/// <returns></returns>
		public static Instances ToInstances(this IDataSet dataSet)
		{
			return ToInstances(dataSet, false);
		}

		/// <summary>
		/// Converts the specified data set into an equivalent 
		/// <see cref="Instances"/> object.
		/// </summary>
		/// <param name="dataSet"></param>
		/// <param name="ignoreLabels">If true, labels are not converted to attributes in the Instances result.</param>
		/// <returns></returns>
		public static Instances ToInstances(this IDataSet dataSet, bool ignoreLabels)
		{
			FastVector attributes = new FastVector(dataSet.Attributes.Count+1);

			if (!ignoreLabels)
			{
				//Create an attribute to hold labels.
				attributes.addElement(new Attribute("label", (FastVector) null));
			}

			//Used for discrete and set attributes.
			FastVector values;

			foreach (DataAttribute attribute in dataSet.Attributes)
			{
				Attribute wekaAttribute;

				switch (attribute.Type)
				{
					//Weka doesn't support discrete/ordinal attributes, so 
					//they are converted to continuous attributes. 
					case AttributeType.Discrete:
					case AttributeType.Continuous:
						wekaAttribute = new Attribute(attribute.Name);
						break;
					case AttributeType.Set:
						values = new FastVector(attribute.PossibleValues.Length);
						Array.ForEach(attribute.PossibleValues, values.addElement);
						wekaAttribute = new Attribute(attribute.Name, values);
						break;
					default:
						throw new NotImplementedException("Attribute type is not supported: " + attribute.Type);
				}

				attributes.addElement(wekaAttribute);
			}

			Instances instances = new Instances("", attributes, dataSet.Instances.Count);

			foreach (Instance instance in dataSet.Instances)
			{
				weka.core.Instance wekaInstance = new weka.core.Instance(dataSet.Attributes.Count+1);
				wekaInstance.setDataset(instances);

				if (!ignoreLabels)
				{
					//Set the label
					wekaInstance.setValue(0, instance.Label);
				}

				for (int i = 0; i < instance.Values.Length; i++)
				{
					if (!ignoreLabels)
					{
						//Remembe to shift everything in the weka instance by +1 since 0 is the label.
						wekaInstance.setValue(i + 1, instance.Values[i]);
					}
					else
					{
						wekaInstance.setValue(i, instance.Values[i]);
					}
				}

				instances.add(wekaInstance);
			}

			instances.setClassIndex(dataSet.TargetAttributeIndex);

			return instances;
		}

		/// <summary>
		/// Converts from <see cref="Instances"/> to <see cref="IDataSet"/>.
		/// </summary>
		/// <param name="instances"></param>
		/// <returns></returns>
		public static IDataSet ToDataSet(this Instances instances)
		{
			BasicDataSet dataSet = new BasicDataSet();

			//If an attribute named "label" is found in the dataset, it is 
			//used to populate the label property of instances in the dataset.
			int labelIndex = -1;

			for (int i=0; i < instances.numAttributes(); i++)
			{
				Attribute attribute = instances.attribute(i);

				if (attribute.isNominal())
				{
					string[] possibleValues = attribute.enumerateValues().ToEnumerable().Cast<string>().ToArray();

					dataSet.Attributes.Add(DataAttribute.NewSetAttribute(attribute.name(), possibleValues));
				}
				else if (attribute.isNumeric())
				{
					dataSet.Attributes.Add(DataAttribute.NewContinuousAttribute(attribute.name()));
				}
				else if (attribute.name().ToLowerInvariant() == "label")
				{
					labelIndex = i;
				}
				else
				{
					throw new NotSupportedException("Attribute is not supported: " + attribute.name() + ".");
				}
			}

			dataSet.TargetAttributeIndex = instances.classIndex();

			//Copy the data
			for (int i = 0; i < instances.numInstances(); i++)
			{
				Instance instance = new Instance(new double[dataSet.Attributes.Count]);

				//A separate index is used for weka since the weka values may contain a label.
				int wekaIndex = 0;

				for (int j=0; j < dataSet.Attributes.Count; j++, wekaIndex++)
				{
					//If the current value in the weka instance is a label, extract it,
					//and don't increment j since we didn't add an actual value to the 
					//instance.
					if (wekaIndex == labelIndex)
					{
						instance.Label = instances.instance(i).stringValue(wekaIndex);
						j--;
						continue;
					}

					instance.Values[j] = instances.instance(i).value(wekaIndex);
				}

				dataSet.Instances.Add(instance);
			}

			return dataSet;
		}
	}
}
